#include <cusp/detail/device/arch.h>
#include <cusp/detail/device/common.h>
#include <cusp/transpose.h>
#include <cusp/blas.h>
#include <cusp/device/block/qr.h>
#include <thrust/detail/device/dereference.h>
#include <thrust/pair.h>

#include <cusp/detail/profiler.h>

#include <cusp/print.h>
#include <cusp/io/dense_matrix.h>
#include <cusp/io/matrix_market.h>
#include <iostream>
#include <fstream>

namespace cusp
{
namespace precond
{
namespace detail
{
namespace device
{

/** \brief Compute the sizes of the Ahat matrices.
	The number of columns is exact, but the number of rows is
	an upper bound.

	This must be run as one block with multiple threads.

	\param AtFirstI iterator over row_offsets of A transposed
	\param AtLastI
	\param MtFirstI iterator over row_offsets of M transposed
	\param MtLastI
	\param MtFirstJ iterator over column_indices of M transposed
	\param temp vector of IndexType for temporary storage with the same dimension as At.row_offsets
	\param [out] sizesRowsFirst beginning of array where row sizes will be stored
	\param [out] sizesColumnsFirst beginning of array where column sizes will be stored
*/
template<typename RandomAccessIterator1,
	 typename RandomAccessIterator2,
	 typename RandomAccessIterator3,
	 typename RandomAccessIterator4,
	 typename RandomAccessIterator5>
__global__ void
get_sizes(RandomAccessIterator1 AtFirstI, RandomAccessIterator1 AtLastI,
	  RandomAccessIterator2 MtFirstI, RandomAccessIterator2 MtLastI,
	  RandomAccessIterator3 MtFirstJ,
	  RandomAccessIterator4 temp,
	  RandomAccessIterator5 sizesRowsFirst,
	  RandomAccessIterator5 sizesColumnsFirst)
{
	typedef typename RandomAccessIterator5::value_type SizeType;
	using thrust::detail::device::dereference;

	/* for every column in M get the number of nnz */
	/* MtFirst and MtLast are the beginning and, respectively, the end
	   of the row_offsets of M transposed */
	RandomAccessIterator2 MtRow = MtFirstI + threadIdx.x;
	RandomAccessIterator5 sizes = sizesColumnsFirst + threadIdx.x;
	--MtLastI; /* row_offsets has one entry more than the number of rows */
	for(; MtRow < MtLastI; MtRow += blockDim.x, sizes += blockDim.x)
	{
		dereference(sizes) = dereference(MtRow+1) - dereference(MtRow);
	}

	/* for every column of A get the number of nnz */
	RandomAccessIterator1 AtRow = AtFirstI + threadIdx.x;
	RandomAccessIterator4 tI = temp + threadIdx.x;
	--AtLastI;
	for(; AtRow < AtLastI; AtRow += blockDim.x, tI += blockDim.x)
	{
		dereference(tI) = dereference(AtRow+1) - dereference(AtRow);
	}

	__syncthreads();

	/* for every column i of M */
	/*   for each nnz row j in i (j is column of A) */
	/*     sum the number of nnz rows in column j of A (temp[j]) */
	MtRow = MtFirstI + threadIdx.x;
	sizes = sizesRowsFirst + threadIdx.x;
	tI = temp;
	for(; MtRow < MtLastI; MtRow += blockDim.x, sizes += blockDim.x)
	{
		SizeType v = 0;
		for(RandomAccessIterator3 j = MtFirstJ + dereference(MtRow); j < MtFirstJ + dereference(MtRow+1); ++j)
			v += dereference(tI + dereference(j));
		dereference(sizes) = v;
	}
}

/** \brief Get indices of rows needed for Ahats.
	For every Ahat copy the indices of the rows of A which have nonzero items to
	the array starting at \a RowsIdxFirst.

	RowsOffsets[i] is the offset from RowsIdxFirst where the indices for i-th Ahat matrix start.
*/
template<typename RandomAccessIterator1,
	 typename RandomAccessIterator2,
	 typename RandomAccessIterator3,
	 typename RandomAccessIterator4,
	 typename RandomAccessIterator5,
	 typename RandomAccessIterator6,
	 typename RandomAccessIterator7>
__global__
void get_row_indices(RandomAccessIterator1 AtFirstI, RandomAccessIterator1 AtLastI,
		     RandomAccessIterator2 AtFirstJ,
		     RandomAccessIterator3 MtFirstI, RandomAccessIterator3 MtLastI,
		     RandomAccessIterator4 MtFirstJ,
		     RandomAccessIterator5 RowsOffsets,
		     RandomAccessIterator6 RowsIdxFirst,
		     RandomAccessIterator7 keys)
{
	using thrust::detail::device::dereference;

	typename thrust::iterator_difference<RandomAccessIterator1>::type MtNumRows = MtLastI - MtFirstI;
	if (blockIdx.x < MtNumRows)
	{
		RowsIdxFirst += dereference(RowsOffsets + blockIdx.x);
		keys += dereference(RowsOffsets + blockIdx.x);

		RandomAccessIterator3 MtRow = MtFirstI + blockIdx.x;
		RandomAccessIterator4 MtCol = MtFirstJ + dereference(MtRow);
		RandomAccessIterator4 MtColEnd = MtFirstJ + dereference(MtRow + 1);
		for(; MtCol < MtColEnd; ++MtCol)
		{
			RandomAccessIterator1 AtRow = AtFirstI + dereference(MtCol);
			RandomAccessIterator2 AtCol = AtFirstJ + dereference(AtRow);
			RandomAccessIterator2 AtColEnd = AtFirstJ + dereference(AtRow + 1);
			typename thrust::iterator_difference<RandomAccessIterator1>::type AtCol_numRows = AtColEnd - AtCol;
			if (threadIdx.x < AtCol_numRows)
			{
				dereference(RowsIdxFirst + threadIdx.x) = dereference(AtCol + threadIdx.x);
				dereference(keys + threadIdx.x) = blockIdx.x;
			}
			RowsIdxFirst += AtCol_numRows;
			keys += AtCol_numRows;
		}
	}
}

/** \brief Build Ahat matrices.
	Copy values from A to Ahat matrices based on sparsity pattern M and the rowIndices array.
*/
template<typename RandomAccessIterator1, typename RandomAccessIterator2,
	 typename RandomAccessIterator3, typename RandomAccessIterator4,
	 typename RandomAccessIterator5, typename RandomAccessIterator6,
	 typename RandomAccessIterator7, typename RandomAccessIterator8,
	 typename RandomAccessIterator9, typename RandomAccessIterator10>
__global__
void populate_Ahats(RandomAccessIterator1 AFirstI, RandomAccessIterator2 AFirstJ, RandomAccessIterator3 AValues,
		    RandomAccessIterator4 MtFirstI, RandomAccessIterator5 MtFirstJ,
		    RandomAccessIterator6 sizesRows, RandomAccessIterator6 sizesCols,
		    RandomAccessIterator7 rowIndicesFirst, RandomAccessIterator8 rowIndicesOffsets,
		    RandomAccessIterator9 offsets, RandomAccessIterator9 offsetsLast,
		    RandomAccessIterator10 AhatFirst)
{
	using thrust::detail::device::dereference;
	typedef typename RandomAccessIterator3::value_type ValueType;

	sizesRows += blockIdx.x;
	sizesCols += blockIdx.x;
	MtFirstI  += blockIdx.x;
	rowIndicesOffsets += blockIdx.x;

	/* in case we have more matrices than blocks */
	for (offsets += blockIdx.x; offsets < offsetsLast; offsets += gridDim.x)
	{
		RandomAccessIterator10 Ahat = AhatFirst + dereference(offsets);
		RandomAccessIterator7 rowIndices = rowIndicesFirst + dereference(rowIndicesOffsets);
		size_t m = dereference(sizesRows);
		size_t n = dereference(sizesCols);
		size_t size = m * n;
		for (size_t el = threadIdx.x; el < size ; el += blockDim.x)
		{
			size_t i = el / n; /* row index in Ahat */
			size_t k = dereference(rowIndices + i); /* row in A */
			size_t j = el % n; /* col index in Ahat */
			size_t l = dereference(MtFirstJ + (dereference(MtFirstI) + j)); /* col in A */

			size_t c = dereference(AFirstI + k);
			size_t cStop = dereference(AFirstI + (k+1));
			ValueType value(0);
			for (; c < cStop; ++c)
			{
				size_t ci = dereference(AFirstJ + c);
				if (ci >= l)
				{
					if (ci == l) value = dereference(AValues + c);
					break;
				}
			}
			dereference(Ahat + el) = value;
		}

		/* prepare for next loop */
		sizesRows += gridDim.x;
		sizesCols += gridDim.x;
		MtFirstI  += gridDim.x;
		rowIndicesOffsets += gridDim.x;
	}
}

/** \brief Perform QR decomposition of the matrices in an array.
*/
template<typename RandomAccessIterator1, typename RandomAccessIterator2,
	 typename RandomAccessIterator3,
	 typename RandomAccessIterator4, typename RandomAccessIterator5>
__global__
void qr(RandomAccessIterator1 Ahats, RandomAccessIterator2 AhatOffsets,
		size_t num_Ahats,
		RandomAccessIterator3 sizesRows, RandomAccessIterator3 sizesColumns,
		RandomAccessIterator4 QValues, RandomAccessIterator5 QOffsets)
{
	using thrust::detail::device::dereference;

	for (size_t i = blockIdx.x; i < num_Ahats; i += gridDim.x)
	{
		cusp::device::block::qr(
			Ahats + dereference(AhatOffsets + i),
			dereference(sizesRows + i),
			dereference(sizesColumns + i),
			QValues + dereference(QOffsets + i));
	}
}

/** \brief Backward substitution.
	Do backward substitution to solve mhat = R \ chat.
*/
template<typename RandomAccessIterator1, typename RandomAccessIterator2,
	 typename RandomAccessIterator3, typename RandomAccessIterator4,
	 typename RandomAccessIterator5, typename RandomAccessIterator6,
	 typename RandomAccessIterator7, typename RandomAccessIterator8,
	 typename RandomAccessIterator9, typename RandomAccessIterator10>
__global__
void make_mhats(RandomAccessIterator1  RValues,    RandomAccessIterator2  ROffsets,
		size_t num_matrices,
		RandomAccessIterator3  sizesRows,  RandomAccessIterator4  sizesColumns,
		RandomAccessIterator5  QValues,    RandomAccessIterator6  QOffsets,
		RandomAccessIterator7  rowIndices, RandomAccessIterator8  rowIndicesOffsets,
		RandomAccessIterator9  mhatsFirst, RandomAccessIterator10 mhatOffsets )
{
	using thrust::detail::device::dereference;
	typedef typename RandomAccessIterator1::value_type ValueType;

	/* use only one thread per block */
	if (threadIdx.x != 0) return;

	for (size_t nMat = blockIdx.x; nMat < num_matrices; nMat += gridDim.x)
	{
		/* find out where the k-th column of M appears (if it does) in Ahat */
		size_t num_rows = dereference(sizesRows + nMat);
		size_t k = 0;
		for (RandomAccessIterator7 AhatRow = rowIndices + dereference(rowIndicesOffsets + nMat);
				k < num_rows; ++AhatRow)
		{
			if (dereference(AhatRow) == nMat) break;
			++k;
		}

		/* do back substitution */
		size_t num_cols = dereference(sizesColumns + nMat);
		RandomAccessIterator9 mhat = mhatsFirst + dereference(mhatOffsets + nMat);
		if (k < num_rows)
		{
			RandomAccessIterator1 R = RValues + dereference(ROffsets + nMat);
			RandomAccessIterator5 Q = QValues + dereference(QOffsets + nMat);

			/* for every row in R, starting from last one */
			for (size_t r = num_cols; r > 0; --r)
			{
				size_t i = r - 1;
				ValueType sum = dereference(Q + (i * num_rows + k));

				/* for every element on the row above the diagonal */
				for (size_t j = r; j < num_cols; ++j)
				{
					sum -= dereference(R + (i * num_cols + j)) * dereference(mhat + j);
				}

				dereference(mhat + i) = sum / dereference(R + (i * num_cols + i));
			}
		}
		else
		{
			for (size_t i = 0; i < num_cols; ++i)
			{
				dereference(mhat + i) = ValueType(0);
			}
		}
	}
}

/** \brief Operator multiplying the elements of a tuple.
	This unary operator takes a tuple with two elements as argument and
	returns the product of these two elements.
*/
template<typename ValueType>
struct multiply_pair : public thrust::unary_function<thrust::tuple<ValueType,ValueType>, ValueType>
{
	__host__ __device__
	ValueType operator()(thrust::tuple<ValueType,ValueType> const& pair)
	{
		return thrust::get<0>(pair) * thrust::get<1>(pair);
	}
};

/** \brief SPAI preconditioner implementation.

	\tparam MatrixType  a cusp::csr_format matrix
	\tparam MatrixType2 a cusp::csr_format matrix

	\param A input matrix
	\param [inout] M sparsity pattern and output preconditioner matrix
*/
template<typename MatrixType,
	 typename MatrixType2>
void spai_csr(const MatrixType& A, MatrixType2 &M)
{
	PROFILE_SCOPED();

	typedef typename MatrixType::index_type IndexType;
	typedef typename MatrixType::value_type ValueType;

	//const unsigned int BLOCK_SIZE = 256;
	//const unsigned int MAX_BLOCKS = cusp::detail::device::arch::max_active_blocks(spai_csr_kernel<DevMatrixType>, BLOCK_SIZE, (size_t) 0);
	//const unsigned int WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;
	//const unsigned int num_blocks = 128;

	/* assume that A is square */
	assert(A.num_rows == A.num_cols);

	/* the sparsity pattern must have the same size as A */
	assert(A.num_rows == M.num_rows);
	assert(A.num_cols == M.num_cols);

	/* create transposed of A and Pattern */
	MatrixType At, Mt;
	cusp::transpose(A, At);
	cusp::transpose(M, Mt);

	typedef typename MatrixType::row_offsets_array_type::iterator RowIteratorType;
	typedef typename MatrixType::column_indices_array_type::iterator ColIteratorType;
	const RowIteratorType MtFirstI = Mt.row_offsets.begin();
	const RowIteratorType MtLastI  = Mt.row_offsets.end();
	const ColIteratorType MtFirstJ = Mt.column_indices.begin();
	//const ColIteratorType MtLastJ  = Mt.column_indices.end();
	const RowIteratorType AtFirstI = At.row_offsets.begin();
	const RowIteratorType AtLastI  = At.row_offsets.end();
	typedef typename cusp::array1d<size_t,cusp::device_memory> SizesType;
	SizesType sizesRows(Mt.num_rows);
	SizesType sizesColumns(Mt.num_rows);
	SizesType tempSize(At.num_rows);

	get_sizes<<<1, 512>>>(AtFirstI, AtLastI,
					      MtFirstI, MtLastI,
					      MtFirstJ,
					      tempSize.begin(),
					      sizesRows.begin(), sizesColumns.begin());
	cusp::array1d<IndexType,cusp::host_memory> sizesRowsH = sizesRows;
	cusp::array1d<IndexType,cusp::host_memory> sizesColumnsH = sizesColumns;

	/* make offsets out of number of rows */
	SizesType rowIndicesOffsets(sizesRows.size() + 1);
	thrust::inclusive_scan(sizesRows.begin(), sizesRows.end(), rowIndicesOffsets.begin() + 1);
	rowIndicesOffsets.front() = 0;

	/* allocate space for row indices */
	cusp::array1d<IndexType,cusp::device_memory> rowIndices(rowIndicesOffsets.back());
	cusp::array1d<size_t,cusp::device_memory> keys(rowIndices.size());

	get_row_indices<<<Mt.num_rows, 512>>>(AtFirstI, AtLastI, At.column_indices.begin(),
					     MtFirstI, MtLastI, MtFirstJ,
					     rowIndicesOffsets.begin(),
					     rowIndices.begin(), keys.begin());

	typename cusp::array1d<IndexType,cusp::device_memory>::iterator start, end;
	start = rowIndices.begin();
	for (typename cusp::array1d<IndexType,cusp::host_memory>::iterator r = sizesRowsH.begin();
			r != sizesRowsH.end(); ++r)
	{
		end = start + *r;
		thrust::sort(start, end); /* required by unique */
		*r = thrust::distance(start, thrust::unique(start, end));
		start = end;
	}
	/* copy new sizes to device */
	thrust::copy(sizesRowsH.begin(), sizesRowsH.end(), sizesRows.begin());

	/* allocate space for Ahat matrices */
	SizesType AhatOffsets(Mt.num_rows + 1);
	thrust::inclusive_scan(
		thrust::make_transform_iterator(thrust::make_zip_iterator(thrust::make_tuple(sizesRows.begin(), sizesColumns.begin())),
			multiply_pair<ValueType>()),
		thrust::make_transform_iterator(thrust::make_zip_iterator(thrust::make_tuple(sizesRows.end(), sizesColumns.end())),
			multiply_pair<ValueType>()),
		AhatOffsets.begin() + 1);
	AhatOffsets.front() = 0;
	cusp::array1d<ValueType,cusp::device_memory> Ahats(AhatOffsets.back());

	populate_Ahats<<<Mt.num_rows, 128>>>(A.row_offsets.begin(), A.column_indices.begin(),
			A.values.begin(),
			MtFirstI, MtFirstJ,
			sizesRows.begin(), sizesColumns.begin(),
			rowIndices.begin(), rowIndicesOffsets.begin(),
			AhatOffsets.begin(), AhatOffsets.end() - 1,
			Ahats.begin());

	/* dump Ahat matrices */
	cusp::array1d<ValueType,cusp::host_memory> AhatsH = Ahats;
	typename cusp::array1d<ValueType,cusp::host_memory>::iterator el = AhatsH.begin();
	typename cusp::array1d<IndexType,cusp::host_memory>::iterator rs = sizesRowsH.begin();
	typename cusp::array1d<IndexType,cusp::host_memory>::iterator cs = sizesColumnsH.begin();
	cusp::array1d<IndexType,cusp::host_memory> rowIndicesH = rowIndices;
	typename cusp::array1d<IndexType,cusp::host_memory>::iterator ri = rowIndicesH.begin();
	typename cusp::array1d<IndexType,cusp::host_memory> rowIndicesOffsetsH = rowIndicesOffsets;

	/* allocate Q matrices */
	cusp::array1d<IndexType,cusp::device_memory> QOffsets(AhatOffsets.size());
	thrust::inclusive_scan(
		thrust::make_transform_iterator(sizesRows.begin(), cusp::blas::detail::square<ValueType>()),
		thrust::make_transform_iterator(sizesRows.end(), cusp::blas::detail::square<ValueType>()),
		QOffsets.begin() + 1);
	QOffsets.front() = 0;

	cusp::array1d<ValueType,cusp::device_memory> QValues(QOffsets.back(), ValueType(0));

	/* QR decomposition */
	qr<<<Mt.num_rows,128>>>(Ahats.begin(), AhatOffsets.begin(),
		AhatOffsets.size() - 1,
		sizesRows.begin(), sizesColumns.begin(),
		QValues.begin(), QOffsets.begin());


	cusp::array1d<IndexType,cusp::device_memory> mhatOffsets(AhatOffsets.size());
	thrust::inclusive_scan(sizesColumns.begin(), sizesColumns.end(), mhatOffsets.begin() + 1);
	mhatOffsets.front() = 0;

	thrust::fill(Mt.values.begin(), Mt.values.end(), ValueType(0));
	make_mhats<<<Mt.num_rows,1>>>(Ahats.begin(), AhatOffsets.begin(),
		Mt.num_rows,
		sizesRows.begin(),  sizesColumns.begin(),
		QValues.begin(),    QOffsets.begin(),
		rowIndices.begin(), rowIndicesOffsets.begin(),
		Mt.values.begin(),  mhatOffsets.begin() );

	cusp::transpose(Mt, M);

#ifdef CUSP_SPAI_DUMP
	/* dump Ahat, Q and R matrices to file */

	std::ofstream out("qrs.m");
	out << "QRs_count = " << Mt.num_rows << ";\n";
	out << "rTol = 1e-5;\n";
	el = AhatsH.begin();
	rs = sizesRowsH.begin();
	cs = sizesColumnsH.begin();
	cusp::array1d<ValueType,cusp::host_memory> QValuesH = QValues;
	typename cusp::array1d<ValueType,cusp::host_memory>::iterator qel = QValuesH.begin();
	cusp::array1d<ValueType,cusp::host_memory> RValuesH = Ahats;
	typename cusp::array1d<ValueType,cusp::host_memory>::iterator rel = RValuesH.begin();
	for (size_t k = 0; k < Mt.num_rows; ++k)
	{
		out << " Ahat = [ ";
		cusp::io::write_dense_matrix(cusp::make_array2d_view(*rs, *cs, *cs,
				cusp::make_array1d_view(el, el + *rs * *cs),
				cusp::row_major()),
			out, " ", " ;\n");
		out << " ];\n";
		out << " Q = [ ";
		cusp::io::write_dense_matrix(cusp::make_array2d_view(*rs, *rs, *rs,
				cusp::make_array1d_view(qel, qel + *rs * *rs),
				cusp::column_major()),
			out, " ", " ;\n");
		out << " ];\n";
		out << " R = [ ";
		cusp::io::write_dense_matrix(cusp::make_array2d_view(*rs, *cs, *cs,
				cusp::make_array1d_view(rel, rel + *rs * *cs),
				cusp::row_major()),
			out, " ", " ;\n");
		out << " ];\n";
		el += *rs * *cs;
		qel += *rs * *rs;
		rel += *rs * *cs;
		++rs; ++cs;

		out << "check(Ahat, Q, R, " << k+1 << ", A, M);\n\n";
	}
	out.close();

	cusp::io::write_matrix_market_file(M, "M.mtx");
#endif
}


} // end namespace device
} // end namespace detail
} // end namespace precond
} // end namespace cusp

